package main

import (
	"bufio"
	"bytes"
	"fmt"
	"golang.org/x/text/encoding/simplifiedchinese"
	"golang.org/x/text/transform"
	"io"
	"net"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"time"
)

// 添加 GBK 到 UTF8 的转换函数
func convertGBKToUTF8(s []byte) ([]byte, error) {
	reader := transform.NewReader(bytes.NewReader(s), simplifiedchinese.GBK.NewDecoder())
	d, e := io.ReadAll(reader)
	if e != nil {
		return nil, e
	}
	return d, nil
}

// 添加网络诊断功能
func diagnoseConnection(addr string) {
	host, _, err := net.SplitHostPort(addr)
	if err != nil {
		fmt.Printf("解析地址失败: %v\n", err)
		return
	}

	fmt.Printf("\n开始网络诊断 %s...\n", host)

	// 首先进行DNS解析
	ips, err := net.LookupIP(host)
	if err != nil {
		fmt.Printf("DNS解析失败: %v\n", err)
	} else {
		fmt.Printf("DNS解析结果: %v\n", ips)
	}

	// 根据操作系统执行不同的网络诊断命令
	var cmd *exec.Cmd
	switch runtime.GOOS {
	case "windows":
		cmd = exec.Command("tracert", host)
	case "linux", "darwin":
		cmd = exec.Command("traceroute", host)
	default:
		fmt.Println("不支持的操作系统")
		return
	}

	output, err := cmd.CombinedOutput()
	if err != nil {
		fmt.Printf("执行网络诊断失败: %v\n", err)
		return
	}

	// 在 Windows 系统下进行编码转换
	if runtime.GOOS == "windows" {
		output, err = convertGBKToUTF8(output)
		if err != nil {
			fmt.Printf("编码转换失败: %v\n", err)
			return
		}
	}

	fmt.Printf("\n路由追踪结果:\n%s\n", string(output))
}

func connect(addr string) (net.Conn, error) {
	fmt.Printf("正在连接 %s...\n", addr)

	// 设置连接超时
	dialer := net.Dialer{
		Timeout:   10 * time.Second,
		KeepAlive: 10 * time.Second,
	}

	// 尝试建立连接
	conn, err := dialer.Dial("tcp", addr)
	if err != nil {
		// 连接失败时进行网络诊断
		diagnoseConnection(addr)
		return nil, fmt.Errorf("连接失败: %v", err)
	}

	return conn, nil
}

func handleConnection(conn net.Conn, limit int) {
	// 添加退出通道
	done := make(chan bool)

	// 添加倒计时计数器
	countdown := limit
	ticker := time.NewTicker(1 * time.Second)
	defer ticker.Stop()

	// 启动读取远程数据的协程
	go func() {
		buffer := make([]byte, 1024)
		for {
			n, err := conn.Read(buffer)
			if err != nil {
				if err != io.EOF && !strings.Contains(err.Error(), "use of closed network connection") {
					fmt.Printf("读取错误: %v\n", err)
				}
				done <- true
				return
			}
			fmt.Print(string(buffer[:n]))
		}
	}()

	// 等待超时或手动退出
	for {
		select {
		case <-ticker.C:
			countdown--
			fmt.Printf("\r%d秒后自动关闭连接", countdown)
			if countdown <= 0 {
				fmt.Println("\n连接空闲超时，正在关闭连接...")
				conn.Close()
				return
			}
		case <-done:
			fmt.Println("\n连接已关闭")
			return
		}
	}
}

// 添加域名解析函数
func resolveDomain(addr string) (string, error) {
	host, port, err := net.SplitHostPort(addr)
	if err != nil {
		return "", fmt.Errorf("解析地址失败: %v", err)
	}

	// 检查是否是IP地址
	if net.ParseIP(host) != nil {
		return addr, nil
	}

	// 进行DNS解析
	ips, err := net.LookupIP(host)
	if err != nil {
		return "", fmt.Errorf("DNS解析失败: %v", err)
	}

	// 优先使用IPv4地址
	for _, ip := range ips {
		if ipv4 := ip.To4(); ipv4 != nil {
			return fmt.Sprintf("%s:%s", ipv4.String(), port), nil
		}
	}

	// 如果没有IPv4地址，使用第一个IPv6地址
	if len(ips) > 0 {
		return fmt.Sprintf("[%s]:%s", ips[0].String(), port), nil
	}

	return "", fmt.Errorf("未找到可用的IP地址")
}

func readInput() (string, error) {
	maxRetries := 3
	for i := 0; i < maxRetries; i++ {
		reader := bufio.NewReader(os.Stdin)
		input, err := reader.ReadString('\n')

		if err != nil {
			if err == io.EOF {
				fmt.Println("\n检测到输入结束，正在重试...")
				time.Sleep(100 * time.Millisecond)
				continue
			}
			return "", err
		}

		if input == "" {
			continue
		}

		return input, nil
	}
	return "", fmt.Errorf("达到最大重试次数")
}

func main() {
	limit := 5
	for {
		fmt.Print("Ctrl+C 或 输入exit、quit退出 \n目标地址 (格式: 域名/IP:端口): ")

		input, err := readInput()
		if err != nil {
			fmt.Printf("读取输入失败: %v\n", err)
			continue
		}

		// 去除输入的空白字符
		addr := strings.Replace(strings.TrimSpace(input), "：", ":", 1)
		// 检查是否退出
		if addr == "exit" || addr == "quit" {
			fmt.Println("程序退出")
			return
		}

		//包含空格
		if strings.Contains(addr, " ") {
			addr = strings.ReplaceAll(addr, " ", ":")
		}
		arr := strings.Split(addr, ":")
		if len(arr) != 2 {
			fmt.Println("输入的地址不合法")
			continue
		}
		resolvedAddr, err := resolveDomain(addr)
		if err != nil {
			fmt.Printf("域名解析失败: %v\n", err)
			continue
		}

		if resolvedAddr != addr {
			fmt.Printf("域名 %s 解析到 IP: %s\n", addr, resolvedAddr)
		}

		t1 := time.Now()
		// 尝试建立连接
		conn, err := connect(resolvedAddr)
		if err != nil {
			fmt.Printf("尝试建立连接错误：%v\n", err)
			continue
		}

		fmt.Printf("已连接到 %s (%s) 耗时%s \n", addr, resolvedAddr, time.Since(t1).String())

		// 处理连接
		handleConnection(conn, limit)
	}
}
